import os
import cv2
import torch
import numpy as np
from PIL import Image
from unet import UNet
import torch.nn.functional as F
from torchvision import transforms
from utils.data_loading import BasicDataset


def predict_img(net, full_img, device, scale_factor=0.5, out_threshold=0.5):

    net.eval()
    # np 转 torch 格式
    img = torch.from_numpy(BasicDataset.preprocess(full_img, scale_factor, is_mask=False))
    # 扩展唯维度
    img = img.unsqueeze(0)
    # 使用GPU
    img = img.to(device=device, dtype=torch.float32)

    with torch.no_grad():
        output = net(img)

        if net.n_classes > 1:
            probs = F.softmax(output, dim=1)[0]
        else:
            probs = torch.sigmoid(output)[0]

        tf = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((full_img.size[1], full_img.size[0])),
            transforms.ToTensor()
        ])

        full_mask = tf(probs.cpu()).squeeze()

    if net.n_classes == 1:
        return (full_mask > out_threshold).numpy()
    else:
        return F.one_hot(full_mask.argmax(dim=0), net.n_classes).permute(2, 0, 1).numpy()


def mask_to_image(mask: np.ndarray):

    if mask.ndim == 2:
        return Image.fromarray((mask * 255).astype(np.uint8))
    elif mask.ndim == 3:
        return Image.fromarray((np.argmax(mask, axis=0) * 255 / mask.shape[0]).astype(np.uint8))


# 与类别修改相对应
net = UNet(n_channels=1, n_classes=3, bilinear=False)

# # 使用CPU
# device = torch.device('cpu')
# # 使用GPU
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda')

# 加载权重文件
model = "checkpoints/checkpoint_epoch10.pth"
net.load_state_dict(torch.load(model, map_location=device))

net.to(device=device, dtype=torch.float32)
# 输入输出路径
path = "./images/input/"
path1 = "./images/output/"

# # 测试时间
# e0 = cv2.getTickCount()

l = 0

for i in os.listdir(path):
    # print(i)
    # 测试时间
    e1 = cv2.getTickCount()

    img1 = cv2.imread(path + i)
    img2 = cv2.resize(img1, (500, 500))

    img3 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)

    img4 = cv2.cvtColor(img3, cv2.COLOR_BGR2GRAY)

    img = Image.fromarray(img4)

    # 预测
    mask = predict_img(net=net, full_img=img, device=device)
    #
    result = mask_to_image(mask)
    result.save(path1 + i)
    l += 1

    # 需要测试的代码
    e2 = cv2.getTickCount()
    time_ = (e2 - e1) / cv2.getTickFrequency()
    print("图像识别时长：", time_)

    # time.sleep(0.1)

# # 需要测试的代码
# e3 = cv2.getTickCount()
# time = (e3 - e0) / cv2.getTickFrequency()
# print("平均每张图像识别时长：", time / l)
